import re
from metagpt.actions import Action
from metagpt.logs import logger
import time
from datetime import datetime
import requests

class CorrectCode(Action):
    #! Take as input a C/C++ code, a Rust code, and the warning and error messages from the Rust compiler, and return a corrected Rust code that is a translation of the input C/C++ code. Write the corrected Rust code to a file named `main.rs`.
    
    name: str = "CorrectCode"
    
    ROLE_DESCRIPTION: str = """In this team, you are an expert in Rust programs. Below is a piece of Rust code and compiler, namely rustc's warning and error messages. The Rust code is translated from C++ program by a new learner of Rust. Please modify the Rust code with the help of the compiler's message to make it compile and run successfully, while aligned with the reference C/C++ program. But do not delete the main function. Think step by step, and eventually aggregate your thoughts in an improved version of the code. Response with:
    ```rust
        Your corrected code here
    ```"""
    
    PROMPT_TEMPLATE: str = """{TEAM_DESCRIPTION}\n{ROLE_DESCRIPTION}
    ## Reference C/C++ Program
    {cpp_code}
    ## Input Translated code
    {rust_code}
    ## Warning and error messages from compiler
    {compiler_message}
    ## Your Thoughts and Final Code
    """

    async def run(self, team_desc: str, filenames: str):
        print(filenames)
        (cpp_file, rust_file, message_file) = filenames.split(";")
        with open(message_file, "r") as f:
            compiler_message = f.read()
        with open(cpp_file, "r") as f:
            cpp_code = f.read()
        with open(rust_file, "r") as f:
            rust_code = f.read()
        logger.info(f"Modifying. Messages: {compiler_message=}")
        prompt = self.PROMPT_TEMPLATE.format(TEAM_DESCRIPTION=team_desc, ROLE_DESCRIPTION=self.ROLE_DESCRIPTION, cpp_code=cpp_code, rust_code=rust_code, compiler_message=compiler_message)

        rsp = await self._aask(prompt)

        rust_code = CorrectCode.parse_code(rsp)
        
        with open(rust_file, "w") as f:
            f.write(rust_code)
            
        
        timestamp = time.time()
        # 将秒数转换为日期和时间
        dt_object = datetime.fromtimestamp(timestamp)
        # 将日期和时间转换为字符串
        human_readable_time = dt_object.strftime("%Y-%m-%d %H:%M:%S")
        data = {"time": human_readable_time,"compare":timestamp,"sender":"code_corrector","receiver":"rag_agent","content":rust_code}
        response = requests.post("http://localhost:5400/post_msg",json=data)
        response_data = response.json()
        print(response_data)

        return cpp_file + ";" + rust_file + ";" + message_file
    
    @staticmethod
    def parse_code(rsp):
        pattern = r"```rust(.*)```"
        match = re.search(pattern, rsp, re.DOTALL)
        code_text = match.group(1) if match else rsp
        return code_text
